import random
import json
import requests
import os
from pathlib import Path

def glob_files(directory, ext='.jpg'):
    dir_path = Path(directory)
    # 递归查找所有 .jpg 后缀的文件（忽略大小写，如 .JPG 也会被匹配）
    files = list(dir_path.rglob(f"*{ext}")) + list(dir_path.rglob(f"*{ext.upper()}"))
    return files

def queue_prompt(prompt):
    """ 提交工作流到 ComfyUI API
    :param prompt: 工作流 JSON 字符串
    :return: API 响应 JSON 数据
    """
    p = {"prompt": prompt}

    data = json.dumps(p)
    response = requests.post("http://127.0.0.1:1082/prompt",
                             data=data,
                             headers={'Content-Type': 'application/json'})
    response.raise_for_status()
    return response.json()

class UdfWorkflow:
    def __init__(self, path_json):
        self.load_workflow(path_json)

    def load_workflow(self, path_json: Path):
        """ 从 JSON 文件加载工作流
        :param path_json: 工作流 JSON 文件路径
        """
        with open(path_json) as fp:
            self.prompt = json.loads(fp.read())

    def run(self, input_img: Path):
        pass

class MaskTagWorkflow(UdfWorkflow):
    def run(self, input_img: Path):
        prompt = self.prompt.copy()
        prompt["33"]["inputs"]["seed"] = random.randint(0, 2**32 - 1)

        prompt["4"]["inputs"]["image"] = str(input_img)
        # output_file = os.path.splitext(os.path.basename(file_path))[0] + ".txt"
        output_file = input_img.with_suffix('.txt').name
        # print(">>> output_file: ", output_file)
        prompt["45"]["inputs"]["file"] = output_file

        # 提交工作流到 ComfyUI API
        result = queue_prompt(prompt)
        # print(f">>> result_json: {result}")
        return result


if __name__ == '__main__':
    wf = MaskTagWorkflow('mask-tag.json')
    for file_path in glob_files('../input'):
        result = wf.run(file_path)
